#include "sweSolver.hpp"
#include "sweFunctions.hpp"
#include "hllc.hpp"
#include "geometry.hpp"

namespace hms
{

/** @brief Calculates the fluxes over all edges and subtracts them from the
 * storage field */
template<typename MeshType>
void SWESolver<MeshType>::calculateFluxesImp(){

	#ifdef SOLVER_STORE_FLUXES
	/* at the beginning of each time step, fluxes need to be set to zero */
		OMP_PRAGMA( parallel )
		{
			auto [colBeg, nCols] = ompSegment( this->fluxes().values().cols() );
			this->fluxes().values().middleCols(colBeg, nCols) = 0;
		}
	#endif

	/* water elevation is reconstructed either as a field or when needed */
	#ifdef SWE_STORE_WATER_ELEVATION
		OMP_PRAGMA( parallel )
		{
			auto [colBeg, nCols] = ompSegment( h().values().cols() );
			h().values().middleCols(colBeg, nCols) = 
				z().values().middleCols(colBeg, nCols) +
				d().values().middleCols(colBeg, nCols);
		}
	#endif

	// meshPartAdjacency is only important for MPI computations. 
	// This is not a solution! Should be perceived as a reminder or placeholder.
	Array<bool, 4, 1> meshPartAdjacency {true, true, true, true};

	const auto& settings = this->settings();


	/* We wrap the scheme call in a lambda, because edge traversal functions 
	 * do not expect member functions */
	/* sadly, wrapping the scheme call twice has an overhead that does not
	 * get optimised out, so we're using a macro for now */
	// auto call = [&]( auto scheme ){
	// 	traverseEdges( this->edgeTraversal(), this->mesh(),
	// 		[&]( const MeshBase<MeshType>& mesh, Index e ){
	// 			(this->*scheme)(mesh, e);
	// 		} );
	// };
	using RT = ReconstructionType;
	#define CALL_FLUX_SCHEME(NAME) \
		traverseEdges( \
			settings.edgeTraversal, this->mesh(), \
			[&]( const MeshBase<MeshType>& mesh, Index e ){ NAME(mesh, e); } )

	#define CALL_FLUX_SCHEME_2ND_ORDER(NAME1, NAME2) \
		traverseEdgesSecondOrder( \
			settings.edgeTraversal, this->mesh(), meshPartAdjacency, \
			[&]( const MeshBase<MeshType>& mesh, Index e){ NAME1(mesh, e); }, \
			[&]( const MeshBase<MeshType>& mesh, Index e){ NAME2(mesh, e); } )

	/* using a comma definition allows us to use a comma as a delimiter between
	 * template arguments. Otherwise the macro-call interprets a comma as a
	 * delimiter for macro arguments.
	 * https://stackoverflow.com/questions/13842468/comma-in-c-c-macro */
	#define COMMA ,

	using FS = FluxScheme;

	switch( settings.fluxScheme ){
		
		case FS::HLLC:
			if ( settings.useMuscl ) {
				if( settings.useAtomicWrite )
					CALL_FLUX_SCHEME_2ND_ORDER(
						hllcScheme<RT::central COMMA true>,
						hllcScheme<RT::secondOrder COMMA true>);
				else
					CALL_FLUX_SCHEME_2ND_ORDER(
						hllcScheme<RT::central COMMA false>,
						hllcScheme<RT::secondOrder COMMA false>);
			}
			else {
				if( settings.useAtomicWrite )
					CALL_FLUX_SCHEME(hllcScheme<RT::firstOrder COMMA true>);
				else
					CALL_FLUX_SCHEME(hllcScheme<RT::firstOrder COMMA false>);
			}
			break;

		case FS::FullyUpwind:
			if ( settings.useAtomicWrite )
				CALL_FLUX_SCHEME(fullyUpwind<true>);
			else
				CALL_FLUX_SCHEME(fullyUpwind<false>);
			break;
		default:
			assert( false &&
				"Flux scheme not registered in calculateFluxesImp()." );
			UNREACHABLE;
	}

	#undef CALL_FLUX_SCHEME
	#undef CALL_HLLC_2ND_ORDER
	#undef COMMA

}

template<typename MeshType>
template<bool atomicWrite>
void SWESolver<MeshType>::registerFlux(
	Ref<Array3s> fluxL,
	Ref<Array3s> fluxR,
	const typename MeshBase<MeshType>::EdgeType& edge,
	const Array2i& lr, // cell indices at edge
	const MeshBase<MeshType>& mesh
){
	fluxL *= this->time().step() * edge.length() / mesh.cell(lr[0]).area();
	fluxR *= this->time().step() * edge.length() / mesh.cell(lr[1]).area();

	#ifdef SOLVER_STORE_FLUXES
		#define FLUX_FIELD this->fluxes()
	#else
		#define FLUX_FIELD q()
	#endif

	#ifdef _OPENMP
	if constexpr (atomicWrite){
		for ( int i = 0; i<3; ++i ){ 
			#pragma omp atomic update
			FLUX_FIELD.values()(i, lr[0]) -= fluxL[i];
			#pragma omp atomic update
			FLUX_FIELD.values()(i, lr[1]) += fluxR[i];
		}
	} else
	#endif
	{
		FLUX_FIELD.values().col(lr[0]) -= fluxL;
		FLUX_FIELD.values().col(lr[1]) += fluxR;
	}

	#undef FLUX_FIELD
};

template<typename MeshType>
template<bool atomicWrite>
void SWESolver<MeshType>::registerFlux(
	Ref<Array3s> flux,
	const typename MeshBase<MeshType>::EdgeType& edge,
	const Array2i& lr, // cell indices at edge
	const MeshBase<MeshType>& mesh
){
	registerFlux<atomicWrite>(flux, flux, edge, lr, mesh);
};

#ifdef SOLVER_TIME_FLUXES
template<typename MeshType>
auto SWESolver<MeshType>::execTimesFlux() -> std::array<double, 8>& {
	return m_execTimesFlux;
}
#endif
template class SWESolver<Mesh>;
template class SWESolver<UniMesh>;
template class SWESolver<RectMesh>;
template class SWESolver<StructMesh>;

} // end namespace hms
